#库的导入
import numpy as np
import matplotlib.pyplot as plt
import heapq

# 待求解问题，求解问题为求最小值
def function(x):
    # x = model_training_func.loss  # 修改为实际的输入变量
    y1 = 0
    for i in range(len(x)-1):
        y2 = 100*((x[i+1] - x[i]**2)**2) + (x[i] - 1)**2
        y1 = y1 + y2
    y = abs(0 - y1)
    y *= 0.001
    return y

class GWO:
    def __init__(self, m=30, imax=30, dimen=3, rangelow=-10, rangehigh=10, amax=2):
        self.m = m
        self.imax = imax
        self.dimen = dimen
        self.rangelow = rangelow
        self.rangehigh = rangehigh
        self.amax = amax
        self.pop = np.zeros((m, dimen))
        self.pop_fitness = np.zeros(m)

    def initialize_population(self):
        for j in range(self.m):
            self.pop[j] = np.random.uniform(low=self.rangelow, high=self.rangehigh, size=(1, self.dimen))
            self.pop_fitness[j] = function(self.pop[j])

    def run(self):
        allbestpop, allbestfit = self.pop[self.pop_fitness.argmin()].copy(), self.pop_fitness.min()
        Xalpha = self.pop[self.pop_fitness.argmin()]
        Xbeta = self.pop[np.argsort(self.pop_fitness)[1]]
        Xdelta = self.pop[np.argsort(self.pop_fitness)[2]]

        his_bestfit = np.zeros(self.imax)

        for i in range(self.imax):
            print("The iteration is:", i + 1)
            iratio = i / self.imax
            a = self.amax * (1 - iratio)

            for j in range(self.m):
                C1 = 2 * np.random.rand()
                Dalpha = np.abs(C1 * Xalpha - self.pop[j])
                A1 = 2 * a * np.random.rand() - a
                X1 = Xalpha - A1 * Dalpha

                C2 = 2 * np.random.rand()
                Dbeta = np.abs(C2 * Xbeta - self.pop[j])
                A2 = 2 * a * np.random.rand() - a
                X2 = Xbeta - A2 * Dbeta
                
                C3 = 2 * np.random.rand()
                Ddelta = np.abs(C3 * Xdelta - self.pop[j])
                A3 = 2 * a * np.random.rand() - a
                X3 = Xdelta - A3 * Ddelta
                

                self.pop[j] = (X1 + X2 + X3) / 3
                self.pop_fitness[j] = function(self.pop[j])

            if self.pop_fitness.min() < allbestfit:
                allbestfit = self.pop_fitness.min()
                allbestpop = self.pop[self.pop_fitness.argmin()].copy()

            his_bestfit[i] = allbestfit
            print("The best fitness is:", allbestfit)

        print("After iteration, the best pop is:", allbestpop)
        print("After iteration, the best fitness is:", "%e" % allbestfit)

        mean = np.sum(self.pop_fitness) /self.m
        std = np.std(self.pop_fitness)
        
        print("After iteration, the mean fitness of the swarm is:", "%e" % mean)
        print("After iteration, the std fitness of the swarm is:", "%e" % std)
        
        return allbestpop

if __name__ == "__main__":
    gwo = GWO()
    gwo.initialize_population()
    gwo.run()

# 将结果进行绘图
#plt.figure(figsize=(8, 4))
#plt.title('The change of best fitness',fontdict={'weight':'normal','size': 30})
#x=range(1,101,1)
#plt.plot(x,his_bestfit,color="red",label="GWO",linewidth=3.0, linestyle="-")
#plt.tick_params(labelsize=25)
#plt.xlim(0,101)
#plt.yscale("log")
#plt.xlabel("Epoch",fontdict={'weight':'normal','size': 30})
#plt.ylabel("Fitness value",fontdict={'weight':'normal','size': 30})
#plt.xticks(range(0,101,10))
#plt.legend(loc="upper right",prop={'size':20})
#plt.savefig("GWO.png")
#plt.show()
